"""
PySCeS - Core2
Author: Brett G. Olivier (brettolivier@gmail.com)
Copyright (C) Brett. G. Olivier, Stellenbosch, 2007.
The source code contained in this file is not licensed for use
or redistribution without the author's prior and express permission.
"""

__version__ = "0.1"

__doc__ = '''PySCeS NewCore module: useful for exploring model component relations'''

import os
CurrentDirectory = os.path.dirname(os.path.abspath(os.sys.argv[0]))

import numpy
##  import scipy

from InfixParser import MyInfixParser

InfixParser = MyInfixParser()
InfixParser.buildlexer()
InfixParser.buildparser(debug=0, debugfile='infix.dbg', tabmodule='infix_tabmodule')
InfixParser.setNameStr('self.', '()')

class MapList(list):
    def __init__(self, *args):
        list.__init__(self,*args)
        
    def asSet(self):
        return set(self.__getslice__(0, self.__len__()))

class NewCoreBase(object):
    name = None

    def getName(self):
        return self.name

    def setName(self,name):
        self.name = name

    def get(self, attr):
        """Return an attribute whose name is str(attr)"""
        try:
            return self.__getattribute__(attr)
        except:
            print "%s is not an attribute of this instance" % attr
            return None

class NumberBase(NewCoreBase):
    value = None
    value_initial = None
    
    def __call__(self):
        return self.value
    
    def getValue(self):
        return self.value
        
    def setValue(self, v):
        self.value = v

class Species(NumberBase):
    subs = None
    prods = None
    mods = None
    fixed = False
    compartment = None

    def __init__(self, name, value):
        self.setName(name)
        self.value = value
        self.value_initial = value
        self.subs = []
        self.prods = []
        self.mods = []

    def setCompartment(self, c):
        self.compartment = c

    def setSubstrate(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.subs.append(reaction)

    def setProduct(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.prods.append(reaction)

    def setModifier(self, reaction):
        self.__setattr__(reaction.name, reaction)
        self.mods.append(reaction)

    def isSubstrateOf(self):
        return MapList([r.name for r in self.subs])

    def isProductOf(self):
        return MapList([r.name for r in self.prods])

    def isModifierOf(self):
        return MapList([r.name for r in self.mods])

    def isReagentOf(self):
        return MapList(self.isSubstrateOf() + self.isProductOf())

class Reaction(NewCoreBase):
    modifiers = None
    substrates = None
    products = None
    stoichiometry = None
    parameters = None
    reversible = True
    formula = None
    code_string = None
    rate = None
    xcode = None
    _names = None

    def __call__(self):
        ##  exec(self.code_string)
        exec(self.xcode)
        return self.rate

    def __init__(self, name):
        self.setName(name)
        self.modifiers = []
        self.substrates = []
        self.products = []
        self.stoichiometry = {}
        self.parameters = []

    def addSubstrate(self, species):
        self.__setattr__(species.name, species)
        self.substrates.append(species)

    def addProduct(self, species):
        self.__setattr__(species.name, species)
        self.products.append(species)

    def addModifier(self, species):
        self.__setattr__(species.name, species)
        self.modifiers.append(species)

    def addFormula(self, formula):
        self.formula = formula
        InfixParser.parse(formula)
        self._names = InfixParser.names
        #print self._names
        self.code_string = 'self.rate=%s' % InfixParser.output
        self.xcode = compile(self.code_string, '<string>', 'exec')

    def addParameter(self, par):
        self.__setattr__(par.name, par)
        self.parameters.append(par)        

    def hasProducts(self, t=type):
        return MapList([p.name for p in self.products])

    def hasSubstrates(self):
        return MapList([s.name for s in self.substrates])

    def hasModifiers(self):
        return MapList([m.name for m in self.modifiers])

    def hasParameters(self):
        return MapList([p.name for p in self.parameters])

    def hasReagents(self):
        return MapList(self.hasSubstrates() + self.hasProducts())


class Parameter(NumberBase):
    association = None
    
    def __init__(self, name, value):
        self.name = name
        self.value = value
        self.value_initial = value
        self.association = []

    def setAssociation(self, reac):
        self.association.append(reac)
        self.__setattr__(reac.name, reac)

    def isParameterOf(self):
        return MapList([a.name for a in self.association])

class AssignmentRule(Parameter):
    formula = None
    code_string = None
    _names = None
    
    def __init__(self, name, value):
        Parameter.__init__(self, name, value)

    def __call__(self):
        exec(self.xcode)
        ##  exec(self.code_string)
        return self.value

    def addFormula(self, formula):
        self.formula = formula
        InfixParser.parse(formula)
        self.code_string = 'self.value=%s' % InfixParser.output
        self._names = InfixParser.names
        self.xcode = compile(self.code_string, '<string>', 'exec')
        
    def addModelAttr(self, obj):
        self.__setattr__(obj.name, obj)

class ODE(NewCoreBase):
    sdot = None
    value = None
    coefficients = None
    reactions = None
    independent = None
    ode_terms = None
    formula = ''
    formula_alt = ''
    code_string = 'self.value='
    code_string_alt = 'sdot='
    
    
    def __init__(self, species, independent=True):
        self.sdot = species
        self.name = 'ODE_'+species.name
        self.reactions = []
        self.coefficients = []
        self.ode_terms = []
        self.independent = independent
                
    def __call__(self):
        exec(self.code_string)
        return self.value

    def addReaction(self, reaction, coefficient):
        self.reactions.append(reaction)
        self.coefficients.append(coefficient)
        if coefficient > 0.0:
            if coefficient == 1.0:
                term = '+self.%s() ' % (reaction.name)
                aterm = '+(%s) ' % (reaction.code_string.replace('self.rate=','')) 
                fterm = '+%s' % (reaction.name)
                afterm = '+ (%s) ' % (reaction.formula) 
            else:
                term = '+%f*self.%s() ' % (abs(coefficient), reaction.name)
                aterm = '+%f*(%s) ' % (abs(coefficient), reaction.code_string.replace('self.rate=','')) 
                fterm = '+%f*%s' % (abs(coefficient), reaction.name)
                afterm = '+ %f*(%s) ' % (abs(coefficient), reaction.formula) 
        else:
            if coefficient == -1.0:
                term = '-self.%s() ' % (reaction.name)
                aterm = '-(%s) ' % (reaction.code_string.replace('self.rate=','')) 
                fterm = '-%s' % (reaction.name)
                afterm = '- (%s) ' % (reaction.formula) 
            else:
                term = '-%f*self.%s() ' % (abs(coefficient), reaction.name)
                aterm = '-%f*(%s) ' % (abs(coefficient), reaction.code_string.replace('self.rate=','')) 
                fterm = '-%f*%s' % (abs(coefficient), reaction.name)
                afterm = '- %f*(%s) ' % (abs(coefficient), reaction.formula) 
        self.ode_terms.append(term)
        self.code_string += term
        self.code_string_alt += aterm
        self.formula += fterm
        self.formula_alt += afterm
        self.__setattr__(reaction.name, reaction)
        ##  for obj in reaction.parameters+reaction.substrates+reaction.products+reaction.modifiers:
            ##  self.__setattr__(obj.name, obj)
    
    def hasReactions(self):
        return MapList([r.name for r in self.reactions])
        
    def getFormula(self):
        return self.code_string
        
    def getGlobalFormula(self):
        return self.code_string_alt
        
    
class StructMatrix(NewCoreBase):
    """
    This class is specifically designed to store structural matrix information
    give it an array and row/col index permutations it can generate its own
    row/col labels given the label src. It might become a more generic class
    but then row/col labels must be externally assigned -- 50/50 atm.
    """
    
    array = None
    ridx = None
    cidx = None
    row = None
    col = None

    def __init__(self, array, ridx, cidx, row=None, col=None):
        """
        Instantiate with array and matching row/col index arrays, optional label arrays
        """
        self.array = array
        self.ridx = ridx
        self.cidx = cidx
        self.row = row
        self.col = col
    
    def __call__(self):
        return self.array
        
    def getRowsByIdx(self, *args):
        """Return the rows referenced by index (1,3,5)"""
        return self.array.take(args, axis=0)
        
    def getColsByIdx(self, *args):
        """Return the columns referenced by index (1,3,5)"""
        return self.array.take(args, axis=1)

    def setRow(self, src):
        """
        Assuming that the row index array is a permutation (full/subset)
        of a source label array by supplying that source to setRow it 
        maps the row labels to ridx and creates self.row (row label list)
        """
        self.row = [src[r] for r in self.ridx]

    def setCol(self, src):
        """
        Assuming that the col index array is a permutation (full/subset)
        of a source label array by supplying that src to setCol 
        maps the row labels to cidx and creates self.col (col label list)
        """
        
        self.col = [src[c] for c in self.cidx]

    def getRowsByName(self, *args):
        """Return the rows referenced by label ('s','x','d')"""
        assert self.row != None, "\nI need row labels"
        try:
            return self.array.take([self.row.index(l) for l in args], axis=0)
        except Exception, ex:
            print ex
            print "\nValid row labels are: %s" % self.row
            return None
        
    def getColsByName(self, *args):
        """Return the columns referenced by label ('s','x','d')"""
        assert self.col != None, "\nI need column labels"
        try:
            return self.array.take([self.col.index(l) for l in args], axis=1)
        except Exception, ex:
            print ex
            print "Valid column labels are: %s" % self.col
            return None

    def getLabels(self, axis='all'):
        """Return the matrix labels ([rows],[cols]) where axis='row'/'col'/'all'"""
        if axis == 'row': return self.row
        elif axis == 'col': return self.col
        else: return self.row, self.col

    def getIndexes(self, axis='all'):
        """Return the matrix indexes ([rows],[cols]) where axis='row'/'col'/'all'"""
        if axis == 'row': return self.ridx
        elif axis == 'col': return self.cidx
        else: return self.ridx, self.cidx

class NewCore(NewCoreBase):
    ndic = None
    reactions = None
    species = None
    species_variable = None
    __model__ = None
    __InitStrings__ = None
    __InitDict__ = None
    __not_inited__ = None
    global_parameters = None
    __parameter_store__ = None
    forcing_functions = None
    __forcing_function_dic__ = None
    __AssignmentRuleClass__ = None
    # new
    compartments = None
    description = "Pysces Core 2: pysXgen"
    global_units = None
    stoichiometric_matrix = None
    struct = None
    ODEs = None

    def __init__(self, model):
        self.setName(model.ModelFile[:-4])
        self.ndic = model._PysMod__NetworkDict
        self.__model__ = model
        self.__InitStrings__ = [s.replace('self.','') for s in self.__model__._PysMod__InitStrings]
        self.__InitDict__ = {}
        # get initialisation strings
        for ii in self.__InitStrings__:
            l,r = ii.split('=')
            self.__InitDict__.setdefault(self.__cleanString__(l), float(self.__cleanString__(r)))
        # set to current model value
        for k in self.__InitDict__.keys():
            self.__InitDict__[k] = getattr(self.__model__, k) 
        self.global_parameters = []
        self.__parameter_store__ = []
        self.__not_inited__ = []
        self.__function_forced_str__ = model._Function_forced
        self.forcing_functions = []
        self.__forcing_function_dic__ = {}
        self.compartments = [('Cell', 1.0)]
        self.global_units = {'mole':(1, 0),
                             'litre':(1, 0),
                             'second':(1, 0)}
        
        try:
            self.parseForcingFunctions()
        except Exception, ex:
            print 'No forcing functions (%s).' % ex
        self.addSpecies()
        self.addReactions()
        self.generateMappings()
        self.setAssignmentRules()
        assert len(self.__not_inited__) < 1, "\nERROR: Uninitialised parameters: %s" % self.__not_inited__
        
    def __cleanString__(self,s):
        s = s.lstrip()
        s = s.rstrip()
        return s

    def parseForcingFunctions(self):
        self.__forcing_function_dic__ = {}
        ff = self.__function_forced_str__.split('\n')
        for f in ff:
            if f != '':
                f =  f.split('=')
                f[0] = f[0].replace('self.','')
                f[1] = f[1].replace('self.','')
                self.__forcing_function_dic__.setdefault(self.__cleanString__(f[0]), self.__cleanString__(f[1]))

    def setDescription(self, txt):
        self.description = str(txt)
        
    def getDescription(self):
        return str(self.description)

    def setGlobalUnits(self, **kwargs):
        for un in kwargs.keys():
            self.global_units[un] = (kwargs[un][0], kwargs[un][1])
            print "Modified \"%s\" to be %i*%s*10**%i" % (un, kwargs[un][0], un, kwargs[un][1])

    def getGlobalUnits(self):
        return self.global_units

    def addOneSpecies(self, species, value, fix=False, comp=None):
        if comp == None: comp = self.compartments[0][0]
        s = Species(species, value)
        s.setCompartment(comp)        
        if fix: s.fixed = True
        self.__setattr__(species, s)
        self.species.append(s)
        if not fix: self.species_variable.append(s)

    def addOneReaction(self, reaction):
        r = Reaction(reaction)
        r.addFormula(self.ndic[r.name]['RateEq'].replace('self.',''))
        if self.ndic[r.name]['Type'] == 'Irrev': r.reversible = False

        #fxnames = self.hasFixedSpecies()
        processed_parameter = []
        # where parameters are defined `locally' per reaction
        for p in self.ndic[r.name]['Params']:
            p = p.replace('self.','')
            if p not in self.hasGlobalParameters() and p not in self.hasFixedSpecies():
                print "\tAdding parameter %s from networkdict" % p
                self.addParameter(p)
                par = self.__getattribute__(p)
                par.setAssociation(r)
                r.addParameter(par)
                processed_parameter.append(p)
            elif p not in self.hasFixedSpecies():
                print "\tUpdating parameter %s from networkdict" % p
                pidx = self.hasGlobalParameters().index(p)
                self.global_parameters[pidx].setAssociation(r)
                r.addParameter(self.global_parameters[pidx])
                processed_parameter.append(p)
        #print self.hasGlobalParameters()
        # where parameters are not `locally' defined and are extracted from Req (ie from SBML)
        for p in r._names:
            p = p.replace('self.','')
            if p not in self.hasGlobalParameters() and p not in self.hasSpecies():
                print "\tAdding parameter %s from formula" % p
                self.addParameter(p)
                par = self.__getattribute__(p)
                par.setAssociation(r)
                r.addParameter(par)
                processed_parameter.append(p)
            elif p not in processed_parameter and p not in self.hasSpecies():
                print "\tUpdating parameter %s from formula" % p
                pidx = self.hasGlobalParameters().index(p)
                self.global_parameters[pidx].setAssociation(r)
                r.addParameter(self.global_parameters[pidx])
                processed_parameter.append(p)
                
        self.__setattr__(reaction, r)
        self.reactions.append(r)

    def addParameter(self, name):
        if self.__InitDict__.has_key(name):
            par = Parameter(name, self.__InitDict__[name])
        else:
            par = Parameter(name, 0.0)
            if name not in self.__not_inited__: self.__not_inited__.append(name)
        self.global_parameters.append(par)
        self.__setattr__(name, par)

    def addSpecies(self):
        self.species = []
        self.species_variable = []
        for s in self.__model__.species:
            if self.__InitDict__.has_key(s):
                self.addOneSpecies(s, self.__InitDict__[s], fix=False)
            else:
                self.addOneSpecies(s, 0.0, fix=False)
        for s in self.__model__.fixed_species:
            if self.__InitDict__.has_key(s):
                self.addOneSpecies(s, self.__InitDict__[s], fix=True)
            else:
                self.addOneSpecies(s, 0.0, fix=True)


    def addReactions(self):
        self.reactions = []
        for r in self.__model__.reactions:
            self.addOneReaction(r)
        non_parameters = self.hasGlobalParameters()+self.hasSpecies()+self.hasFixedSpecies()
        for k in self.__InitDict__.keys():
            if k not in non_parameters:
                print 'Adding new parameter:', k
                self.addParameter(k)

    def replaceParameterWithRule(self, ar):
        par = self.__getattribute__(ar.name)
        for r in par.association:
            ar.setAssociation(r)
            setattr(r, ar.name, ar)
            r.parameters[r.hasParameters().index(ar.name)] = ar
        self.global_parameters[self.hasGlobalParameters().index(ar.name)] = ar
        self.__setattr__(ar.name, ar)

    def setAssignmentRules(self):
        for p in self.global_parameters:
            #print p.name
            if self.__forcing_function_dic__.has_key(p.name):
                print 'Assigning: %s = %s' % (p.name, self.__forcing_function_dic__[p.name])
                p2 = AssignmentRule(p.name, None)
                self.replaceParameterWithRule(p2)
                assert isinstance(p2, AssignmentRule), "\nHappy assertion error"
                #print type(p2)
                p2.addFormula(self.__forcing_function_dic__[p.name])
                print p2._names
                for n in p2._names:
                    p2.addModelAttr(self.__getattribute__(n))
                self.forcing_functions.append(p2)
                if p2.name in self.__not_inited__:
                    self.__not_inited__.pop(self.__not_inited__.index(p.name))

    def generateMappings(self):
        for reac in self.reactions:
            for reag in self.ndic[reac.name]['Reagents']:
                if self.ndic[reac.name]['Reagents'][reag] < 0.0:
                    reac.addSubstrate(self.__getattribute__(reag.replace('self.','')))
                    self.__getattribute__(reag.replace('self.','')).setSubstrate(self.__getattribute__(reac.name))
                else:
                    reac.addProduct(self.__getattribute__(reag.replace('self.','')))
                    self.__getattribute__(reag.replace('self.','')).setProduct(self.__getattribute__(reac.name))
                reac.stoichiometry.setdefault(reag.replace('self.',''), self.ndic[reac.name]['Reagents'][reag])
            for mod in self.ndic[reac.name]['Modifiers']:
                reac.addModifier(self.__getattribute__(mod.replace('self.','')))
                self.__getattribute__(mod.replace('self.','')).setModifier(self.__getattribute__(reac.name))

    def setStoichiometricMatrix(self):
        vspec = self.hasVariableSpecies()
        react = self.hasReactions()
        nm = numpy.zeros((len(vspec), len(react)),'d')
        for sp in vspec:
            for r in self.get(sp).isReagentOf():
                nm[vspec.index(sp)][react.index(r)] = self.get(r).stoichiometry[sp]
            # this is if absolute stoichiometry value is used
            ##  for r in self.get(sp).isSubstrateOf():
                ##  nm[vspec.index(sp)][react.index(r)] = abs(self.get(r).stoichiometry[sp])
            ##  for r in self.get(sp).isProductOf():
                ##  nm[vspec.index(sp)][react.index(r)] = -abs(self.get(r).stoichiometry[sp])
        self.stoichiometric_matrix = StructMatrix(nm, range(len(vspec)), range(len(react)))
        self.stoichiometric_matrix.setRow(vspec)
        self.stoichiometric_matrix.setCol(react)

    def addODEs(self):
        self.ODEs = []
        for varspec in self.stoichiometric_matrix.row:
            if self.struct != None:
                if varspec not in self.struct.Nr.row:
                    print 'Creating dependent ODE_%s' % varspec
                    ode = ODE(self.get(varspec), independent=False)
                else:
                    print 'Creating independent ODE_%s' % varspec
                    ode = ODE(self.get(varspec), independent=True)
            else:
                print 'Creating independent* ODE_%s (*assumed - no structural information available)' % varspec
                ode = ODE(self.get(varspec), independent=True)
            mrow = self.stoichiometric_matrix.getRowsByName(varspec)
            for e in range(len(mrow[0])):
                if mrow[0,e] != 0.0:
                    print '\tAdding term: %s*%s' % (mrow[0,e], self.stoichiometric_matrix.col[e])
                    ode.addReaction(self.get(self.stoichiometric_matrix.col[e]), mrow[0,e])
            self.__setattr__(ode.name, ode)
            self.ODEs.append(ode)
            self.__setattr__('xcode_'+ode.name, compile(ode.getGlobalFormula(), '<string>', 'exec'))
            
    def hasODEs(self):
        return MapList([o.name for o in self.ODEs])

    def evalODEs(self, odes):
        return [v() for v in odes]
        
    def evalXcode(self, ode):
        exec(self.__getattribute__('xcode_'+ode.name))
        return sdot

    def hasReactions(self):
        return MapList([r.name for r in self.reactions])

    def hasSpecies(self):
        return MapList([s.name for s in self.species])

    def hasFixedSpecies(self):
        return MapList([s.name for s in self.species if s.fixed])

    def hasVariableSpecies(self):
        return MapList([s.name for s in self.species if not s.fixed])

    def findReactionsThatIncludeAllSpecifiedReagents(self, *args):
        assert len(args) > 1, '\nNeed two or more species for this one!'
        setlist = [self.__getattribute__(s).isReagentOf().asSet() for s in args]
        isect = setlist[0]
        for s in setlist:
            isect.intersection_update(s)
        return MapList(isect)

    def hasGlobalParameters(self):
        return MapList(p.name for p in self.global_parameters)

##  if __psyco_active__:
    ##  psyco.bind(NewCoreBase)
    ##  psyco.bind(NumberBase)
    ##  psyco.bind(Species)
    ##  psyco.bind(Parameter)
    ##  psyco.bind(AssignmentRule)
    ##  psyco.bind(Reaction)
    ##  psyco.bind(ODE)
    ##  psyco.bind(NewCore)

            

if __name__ == '__main__':
    
    print "\n*****\nTESTING NEWCORE\n*****\n"


    import pysces
    M = pysces.model('isola_thesis')
    M.doLoad()

    print '\nModel', M.ModelFile
    print '============='
    modmap = NewCore(M)

    print 'Reactions\n', modmap.hasReactions()
    print 'Species\n', modmap.hasSpecies()
    print 'FixedSpecies\n', modmap.hasFixedSpecies()
    print ' '
    print 'R1 has reagents\n', modmap.R1.hasReagents()
    print 'R1 has sub\n', modmap.R1.hasSubstrates()
    print 'R1 has prod\n', modmap.R1.hasProducts()
    print 'R1 has mod\n', modmap.R1.hasModifiers()
    print ' '
    print 'P is reagent\n', modmap.P.isReagentOf()
    print 'P is sub\n', modmap.P.isSubstrateOf()
    print 'P is prod\n', modmap.P.isProductOf()
    print 'P is mod\n', modmap.P.isModifierOf()
    print ' '
    print 'R2 stoich\n', modmap.R2.stoichiometry
    print ' '
    print 'findReactionsThatIncludeAllSpecifiedReagents(A, B):', modmap.findReactionsThatIncludeAllSpecifiedReagents('A','B')


    print '\nmodmap.hasGlobalParameters\n', modmap.hasGlobalParameters()

    print '\nParameter associations'
    for p in modmap.global_parameters:
        print '%s.isParameterOf() %s' % (p.name, p.isParameterOf())

    print "\n*****\nTESTING INFIXPARSER\n*****\n"

    # Build the lexer and try it out
    ##  m = MyInfixLexer()
    m = MyInfixParser()
    m.buildlexer()           # Build the lexer
    ##  m.testlexer("(K3*K4+s0)**(h)")     # Test it
    m.buildparser(debug=0, debugfile='infix.dbg', tabmodule='infix_tabmodule')
    m.setNameStr('self.', '()')
    print "\ngo!\n"
    
    m.parse("(K3*K4+s0)**(h)")
    print '\nm.input\n', m.input
    print 'm.output\n', m.output
    print 'm.names\n', m.names
    
    m.parse("(2**K3*K4+s0)**(h/5)")
    print '\nm.input\n', m.input
    print 'm.output\n', m.output
    print 'm.names\n', m.names
